package leetcode_2401_2500;

import java.util.HashMap;
import java.util.Map;

public class LeeCode_2488_3 {
    public static void main(String[] args) {
        System.out.println(countSubarrays(new int[]{3, 2, 1, 4, 5}, 4));
        System.out.println(countSubarrays(new int[]{2,3,1}, 3));
        System.out.println(countSubarrays(new int[]{5,19,11,15,13,16,4,6,2,7,10,8,18,20,1,3,17,9,12,14}, 6));
    }
    private static int countSubarrays(int[] nums, int k) {
        int len = nums.length;
        int start = 0;
        int res = 0;
        for (int i = 0; i < len; i++) {
            if (nums[i] == k){
                start = i;
                nums[i] = 0;
                continue;
            }
            nums[i] = nums[i] < k ? -1 : 1;
        }
        Map<Integer, Integer> counts = new HashMap<>();
        int sum = 0;
        for (int i = 0; i < len; i++) {
            sum += nums[i];
            if (i < start){
                counts.put(sum, counts.getOrDefault(sum, 0) + 1);
            }
            if (i >= start){
                int countPre = counts.getOrDefault(sum, 0);
                int countAft = counts.getOrDefault(sum - 1, 0);
                res += countAft + countPre;
            }
        }
        return res;
    }
}
